import os
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path

import cv2
import pandas as pd
from sklearn.model_selection import train_test_split
from tqdm import tqdm

from constants import CELEB_DF
from option import parse_args
from preprocessing.generate_folds import get_real_fake_pairs_for_test, get_real_fake_pairs, cmp_frame_name, cmp

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)

SAMPLE_INTERVALS = [-40, -30, -20, -10, 0, 10, 20, 30, 40]


def check_crops2(ori_dir, fake_dir, frame):
    frame_split = frame.split('_')
    frame_num = int(frame_split[0])
    for interval in SAMPLE_INTERVALS:
        frame_cur = frame_num + interval
        img_name = f'{frame_cur}_{frame_split[1]}'
        ori_img_path = os.path.join(ori_dir, img_name)
        fake_img_path = os.path.join(fake_dir, img_name)
        if not os.path.exists(ori_img_path) or not os.path.exists(fake_img_path):
            return False
    return True


def get_paths(vid, label, root_dir, sample_num):
    ori_vid, fake_vid = vid
    ori_dir = os.path.join(root_dir, 'crops', ori_vid)
    fake_dir = os.path.join(root_dir, 'crops', fake_vid)
    mask_dir = os.path.join(root_dir, 'diffs', fake_vid)

    data = []
    # Some unaligned videos have been removed
    if label == 1 and not os.path.exists(fake_dir):
        return data
    # Some masks may have been removed
    if label == 1 and not os.path.exists(mask_dir):
        return data

    frames = os.listdir(fake_dir)
    frames.sort(key=cmp_frame_name)
    for frame in frames:
        frame_split = str(frame).split('_')
        frame_num = int(frame_split[0])
        if frame_num < 50 or frame_num % 100 != 50 or not check_crops2(ori_dir, fake_dir, frame):
            continue
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if label == 0 and os.path.exists(ori_img_path):
            data.append([ori_img_path, label, ori_vid])
        elif label == 1 and os.path.exists(fake_img_path) and os.path.exists(ori_img_path):
            mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
            if os.path.exists(mask_path):
                data.append([fake_img_path, label, ori_vid])

    return data


def save_folds(args, save_path, pairs_mode, mode='train', sample_num=100):
    data = []
    if mode == 'train_fake':
        ori_ori_pairs = set()
        ori_fake_pairs = pairs_mode
    elif mode == 'train_real':
        ori_ori_pairs = pairs_mode
        ori_fake_pairs = set()
    else:
        ori_ori_pairs = set()
        ori_fake_pairs = set()
        for real, fake in pairs_mode:
            if real == fake:
                ori_ori_pairs.add((real, fake))
            else:
                ori_fake_pairs.add((real, fake))

    ori_ori_pairs = list(ori_ori_pairs)
    ori_fake_pairs = list(ori_fake_pairs)
    ori_ori_pairs.sort(key=cmp)
    ori_fake_pairs.sort(key=cmp)

    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        video = path.parent.name
        file = path.name
        file_split = file.split("_")
        if len(file_split) == 2:
            frame = int(file_split[0])
        else:
            frame = 0
        fold_data.append([video, file, label, ori_vid, frame])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    file_name = '{}_st_{}.csv'.format(CELEB_DF, mode)
    save_file = os.path.join(save_path, file_name)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def real_fake_split(ori_fake_pairs):
    ori_fake_pairs_real = []
    ori_fake_pairs_fake = []
    for pair in ori_fake_pairs:
        if pair[0] == pair[1]:
            ori_fake_pairs_real.append(pair)
        else:
            ori_fake_pairs_fake.append(pair)
    return ori_fake_pairs_real, ori_fake_pairs_fake


def main():
    args = parse_args()
    save_path = '../data/{}'.format(CELEB_DF)
    os.makedirs(save_path, exist_ok=True)
    pairs_test = get_real_fake_pairs_for_test(args.root_dir)
    ori_fake_pairs = get_real_fake_pairs(args.root_dir)
    assert set(ori_fake_pairs).isdisjoint(set(pairs_test))
    pairs_train, pairs_val = train_test_split(ori_fake_pairs, test_size=0.3, random_state=111)
    print(len(pairs_train), len(pairs_val), len(pairs_test))

    sample_num_train = 20
    sample_num_test = 20
    sample_num_val = 10
    save_folds(args, save_path, pairs_train, 'train', sample_num=sample_num_train)
    save_folds(args, save_path, pairs_val, 'val', sample_num=sample_num_val)
    save_folds(args, save_path, pairs_test, 'test', sample_num=sample_num_test)

    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(pairs_train)
    save_folds(args, save_path, ori_fake_pairs_real, 'train_real', sample_num=sample_num_train)
    save_folds(args, save_path, ori_fake_pairs_fake, 'train_fake', sample_num=sample_num_train)


if __name__ == '__main__':
    # Total: 6529
    # Train: 5409
    # Validation: 602
    # Test: 518
    main()
